// quick code to compute the hemisphere mass
#ifndef __mr__
#define __mr__

#include <vector>
#include <cassert>
#include "TLorentzVector.h"
#include <cmath>

using namespace std;

typedef vector<TLorentzVector> vec_list;


double cal_mr(const vec_list& pair)
{
  assert(pair.size()==2);
  
  TLorentzVector ja=pair[0];
  TLorentzVector jb=pair[1];

  /*
  TLorentzVector ja, jb;

  ja.SetPtEtaPhiM(j1.Pt(),j1.Eta(),j1.Phi(),0.0);
  jb.SetPtEtaPhiM(j2.Pt(),j2.Eta(),j2.Phi(),0.0);
  */

  float A = ja.P();
  float B = jb.P();
  float az = ja.Pz();
  float bz = jb.Pz();
  TVector3 jaT, jbT;
  jaT.SetXYZ(ja.Px(),ja.Py(),0.0);
  jbT.SetXYZ(jb.Px(),jb.Py(),0.0);
  float ATBT = (jaT+jbT).Mag2();

  float temp = sqrt((A+B)*(A+B)-(az+bz)*(az+bz)-(jbT.Dot(jbT)-jaT.Dot(jaT))*(jbT.Dot(jbT)-jaT.Dot(jaT))/(jaT+jbT).Mag2());

  float mybeta = (jbT.Dot(jbT)-jaT.Dot(jaT))/sqrt(ATBT*((A+B)*(A+B)-(az+bz)*(az+bz)));

  float mygamma = 1./sqrt(1.-mybeta*mybeta);

  //gamma times MRstar
  temp *= mygamma;

  return temp;
}

//compute the sum of masses for a vector<TLorentzVector> of size two
double vector_mass(const vec_list& pair)
{
  assert(pair.size()==2);
  return pair[0].M() + pair[1].M();
}

//compute two mega-jet with seed
vec_list megajet_seed(const vec_list& input,
		      TLorentzVector& a)
{
  vec_list result(2);
  
  TLorentzVector seed(a.Vect().Unit(), 0);

  //loop over the input
  //we want to maximum the 3 vector dot product of a*set 1 + b*set 2
  for(int i=0; i<input.size(); ++i)
    {
      //we want to minimize the hemisphere mass with respect to the seed
      if(input[i].Dot(seed) > 0)
	result[0]+= input[i];
      else
	result[1]+= input[i];
    }

  //modify the seed
  a=TLorentzVector(result[0].Vect().Unit(),0);
  return result;
}


//iterate a few times and see what happens
vec_list megajet_it(const vec_list& input,
		    TLorentzVector a)
{
  double current_m=1e9;
  vec_list result;

  for(int i=0; i<10; i++)
    {
      vec_list temp_vec =
	megajet_seed(input, a);

      double temp_m = vector_mass(temp_vec);
      //make sure iteration makes sense
      //with some room for rounding errors

      if(temp_m >= current_m + 1e-6)
	cout<<"error: "<<temp_m<<" , "<<current_m<<endl;

      result=temp_vec;

      if(fabs(temp_m-current_m) < 1e-5)
	break;
      else
	current_m = temp_m;
    }

  return result;
}



//compute the two mega-jet given jet input
vec_list megajet(vec_list input)
{
  //compare the results from different iteration
  TLorentzVector j1, j2;
  double temp_m=0;

  //compute the boost
  TLorentzVector boost4;
  for(int i=0; i<input.size(); ++i)
    boost4 += input[i];

  TVector3 boost=boost4.Vect();
  boost *= (-1.0 / boost4.E());

  //now boost everything
  for(int i=0; i<input.size(); ++i)
    input[i].Boost(boost);
 
  vec_list final_result;
  double final_mass = 1e9;

  for(int i=0; i<input.size() && i<4; ++i)
      {
	vec_list result=megajet_it(input, input[i]);
	temp_m=vector_mass(result);

	if(temp_m < final_mass)
	  {
	    final_mass = temp_m;
	    final_result = result;
	  }
	//cout<<"mass is "<<i<<": "<<temp_m<<endl;
      }

  //boost everything back
  boost*=-1.0;
  final_result[0].Boost(boost);
  final_result[1].Boost(boost);
  
  return final_result;
}



//a brute force method to double check
vec_list megajet_brute_old(vec_list input)
{
  if(input.size() < 2)
    cout<<"ERROR: less than two input jets!"<<endl;

  double final_mass = 1e9;
  vec_list final_vec(2);

  //number of combinations
  int n=pow(2, input.size());
  
  for(int i=0; i<n; i++)
    {
      vec_list temp_vec(2);
      int k=i;

      for(int j=0; j<input.size(); j++)
	{
	  //read the bit
	  temp_vec[k%2]+=input[j];
	  k/=2;	  
	}

      //calculate mass for specific combination
      double temp_mass=vector_mass(temp_vec);
      
      //now compare
      if(temp_mass < final_mass)
	{
	  final_mass = temp_mass;
	  final_vec = temp_vec;
	}

    }
  return final_vec;
}




//Takes as imput the collection of all your jets with pT>30 GeV 
vector<TLorentzVector> megajet_brute
(const vector<TLorentzVector>& myjets){

  vector<TLorentzVector> mynewjets;
  TLorentzVector j1, j2;
  bool foundGood = false;

  int N_comb = 1;
  for(int i = 0; i < myjets.size(); i++){
    N_comb *= 2;
  }

  float M_min = 9999999999.0;
  int j_count;
  for(int i = 1; i < N_comb-1; i++){
    TLorentzVector j_temp1, j_temp2;
    int itemp = i;
    j_count = N_comb/2;
    int count = 0;
    while(j_count > 0){
      if(itemp/j_count == 1){
        j_temp1 += myjets[count];
      } else {
        j_temp2 += myjets[count];
      }
      itemp -= j_count*(itemp/j_count);
      j_count /= 2;
      count++;
    }
    float M_temp = j_temp1.M2()+j_temp2.M2();
    if(M_temp < M_min){
      M_min = M_temp;
      j1 = j_temp1;
      j2 = j_temp2;
    }
  }

  // set masses to 0
  j1.SetPtEtaPhiM(j1.Pt(),j1.Eta(),j1.Phi(),0.0);
  j2.SetPtEtaPhiM(j2.Pt(),j2.Eta(),j2.Phi(),0.0);

  if(j2.Pt() > j1.Pt()){
    TLorentzVector temp = j1;
    j1 = j2;
    j2 = temp;
  }

  mynewjets.push_back(j1);
  mynewjets.push_back(j2);
  return mynewjets;

}


#endif
